#include "BTreeNode.h"

using namespace std;


/********************************************************************/
/*                                                                  */
/*                        BTLeafNode Classes                        */
/*                                                                  */
/********************************************************************/

/* BTLeafNode constructor, initialize everything */
BTLeafNode::BTLeafNode() {
    clearBuffer();
    int* keyCount = (int*) &buffer[KEY_COUNT_INDEX];
    *keyCount = 0;
}

/*
 * Clear the main memory.
 * Sets the entire buffer to -1.
 * @return 0 if successful. Return an error code if there is an error.
 */
RC BTLeafNode::clearBuffer() {
  memset(buffer, -1, PageFile::PAGE_SIZE);
  return 0;
}

/*
 * Read the content of the node from the page pid in the PageFile pf.
 * @param pid[IN] the PageId to read
 * @param pf[IN] PageFile to read from
 * @return 0 if successful. Return an error code if there is an error.
 */
RC BTLeafNode::read(PageId pid, const PageFile& pf) {
  return pf.read(pid, buffer);
}
    
/*
 * Write the content of the node to the page pid in the PageFile pf.
 * @param pid[IN] the PageId to write to
 * @param pf[IN] PageFile to write to
 * @return 0 if successful. Return an error code if there is an error.
 */
RC BTLeafNode::write(PageId pid, PageFile& pf) {
  return pf.write(pid, buffer);
}

/*
 * Return the number of keys stored in the node.
 * @return the number of keys in the node
 */
int BTLeafNode::getKeyCount() { 
    // The number of elements in a leaf node is stored in the last sizeof(int) bytes in the page
  int* keyCount = (int*) &buffer[KEY_COUNT_INDEX];
  return *keyCount;
}

/*
 * Insert a (key, rid) pair to the node.
 * @param key[IN] the key to insert
 * @param rid[IN] the RecordId to insert
 * @return 0 if successful. Return an error code if the node is full.
 */
RC BTLeafNode::insert(int key, const RecordId& rid) {
    // Check if the node is full
  if (getKeyCount() == MAX_ELEMS)
    return RC_NODE_FULL;
    
    // Find position to insert new key
  int insertIndex = 0;
  RC locateError = locate(key, insertIndex);
  insertIndex *= ELEM_SIZE;
  
    // Check for unexpected error
  if (locateError != 0 && locateError != RC_NO_SUCH_RECORD)
    return locateError;
  
    // Shift the existing pairs over without overwriting the next node page id
  memmove(&buffer[insertIndex + ELEM_SIZE], &buffer[insertIndex], PAGE_ID_INDEX - (insertIndex + ELEM_SIZE));
  
    // Insert the new pair
  int* newKey = (int*) &buffer[insertIndex];
  RecordId* newRid = (RecordId*) &buffer[insertIndex + sizeof(int)];
  *newKey = key;
  *newRid = rid;
  
    // Increase keyCount
  int* keyCount = (int*) &buffer[KEY_COUNT_INDEX];
  (*keyCount)++;
  return 0;
}

/*
 * Insert the (key, rid) pair to the node
 * and split the node half and half with sibling.
 * The first key of the sibling node is returned in siblingKey.
 * @param key[IN] the key to insert.
 * @param rid[IN] the RecordId to insert.
 * @param sibling[IN] the sibling node to split with. This node MUST be EMPTY when this function is called.
 * @param siblingKey[OUT] the first key in the sibling node after split.
 * @return 0 if successful. Return an error code if there is an error.
 */
RC BTLeafNode::insertAndSplit(int key, const RecordId& rid, 
                              BTLeafNode& sibling, int& siblingKey) 
{ 
  int keyCount = getKeyCount();
    // Index of the key in the middle, one above if keyCount is even
  int splitCount = keyCount / 2;
  int insertIndex = 0;
  int splitOffset = 0;
  
    // Find the position to insert new key
  locate(key, insertIndex);
    
    // Check if new key goes into current or sibling node
    // Make room for element depending on which node it is in
  if (insertIndex > splitCount) 
    splitOffset = 1;
    
    // Insert everything after split point of current node into sibling
  for (int i = splitCount + splitOffset; i < keyCount; i++) {
    int* keyToCopy = (int*) &buffer[i * ELEM_SIZE];
    RecordId* ridToCopy = (RecordId*) &buffer[i * ELEM_SIZE + sizeof(int)];
    sibling.insert(*keyToCopy, *ridToCopy);
  }
  
    // Clear elements after the split point in current node
  memset(&buffer[(splitCount + splitOffset) * ELEM_SIZE], -1, (keyCount - splitCount - splitOffset) * ELEM_SIZE);
 
    // Insert new key
  if (insertIndex > splitCount)
    sibling.insert(key, rid);
  else
    insert(key, rid);
  
    // Update key count of current node
  int* currentKeyCount = (int*) &buffer[KEY_COUNT_INDEX];
  *currentKeyCount = splitCount + 1;
  
    // Get first key of sibling
  RecordId sibFirstRid;
  sibling.readEntry(0, siblingKey, sibFirstRid);
  
    // Move current next node PageId to sibling
  sibling.setNextNodePtr(getNextNodePtr());
  return 0; 
}

/*
 * Find the entry whose key value is larger than or equal to searchKey
 * and output the eid (entry number) whose key value >= searchKey.
 * Remember that all keys inside a B+tree node should be kept sorted.
 * @param searchKey[IN] the key to search for
 * @param eid[OUT] the entry number that contains a key larger than or equalty to searchKey
 * @return 0 if successful. Return an error code if there is an error.
 */
RC BTLeafNode::locate(int searchKey, int& eid) {
  int keyCount = getKeyCount();
  int i = 0;
  
  for (; i < keyCount; i++) {
    int* currentKey = (int*) &buffer[ELEM_SIZE * i];
      // Stop once a key is found to be greater than searchKey
    if (*currentKey >= searchKey)
      break;
  }
  
  eid = i;
  if (i == keyCount)
    return RC_NO_SUCH_RECORD;
  return 0;
}

/*
 * Read the (key, rid) pair from the eid entry.
 * @param eid[IN] the entry number to read the (key, rid) pair from
 * @param key[OUT] the key from the entry
 * @param rid[OUT] the RecordId from the entry
 * @return 0 if successful. Return an error code if there is an error.
 */
RC BTLeafNode::readEntry(int eid, int& key, RecordId& rid) {
    // eid must be less than the number of pairs 
  if (eid >= getKeyCount())
    return RC_NO_SUCH_RECORD;
    
  int pairIndex = eid * ELEM_SIZE;
  int* storedKey = (int*) &buffer[pairIndex];
  RecordId* storedRid = (RecordId*) &buffer[pairIndex + sizeof(int)];
  key = *storedKey;
  rid = *storedRid;
  return 0; 
}

/*
 * Return the pid of the next sibling node.
 * @return the PageId of the next sibling node 
 */
PageId BTLeafNode::getNextNodePtr() {
  PageId* nextNodePtr = (PageId*) &buffer[PAGE_ID_INDEX];
  return *nextNodePtr;
}

/*
 * Set the pid of the next sibling node.
 * @param pid[IN] the PageId of the next sibling node 
 * @return 0 if successful. Return an error code if there is an error.
 */
RC BTLeafNode::setNextNodePtr(PageId pid) {
  PageId* nextNodePtr = (PageId*) &buffer[PAGE_ID_INDEX];
  *nextNodePtr = pid;
  return 0; 
}


/********************************************************************/
/*                                                                  */
/*                      BTNonLeafNode Classes                       */
/*                                                                  */
/********************************************************************/


/* BTNonLeafNode constructor, initialize everything */
BTNonLeafNode::BTNonLeafNode() {
    clearBuffer();
    int* keyCount = (int*) &buffer[KEY_COUNT_INDEX];
    *keyCount = 0;
}

/* Set first pid of the nonleaf node.
 * Should only be used after insertAndSplit
 */
RC BTNonLeafNode::setFirstPid(PageId pid) {
    PageId* firstPid = (PageId*) &buffer[0];
    *firstPid = pid;
    return 0;
}

/*
 * Clear the main memory.
 * Sets the entire buffer to -1.
 * @return 0 if successful. Return an error code if there is an error.
 */
RC BTNonLeafNode::clearBuffer() {
  memset(buffer, -1, PageFile::PAGE_SIZE);
  return 0;
}

/*
 * Read the content of the node from the page pid in the PageFile pf.
 * @param pid[IN] the PageId to read
 * @param pf[IN] PageFile to read from
 * @return 0 if successful. Return an error code if there is an error.
 */
RC BTNonLeafNode::read(PageId pid, const PageFile& pf) { 
  return pf.read(pid, buffer);
}
    
/*
 * Write the content of the node to the page pid in the PageFile pf.
 * @param pid[IN] the PageId to write to
 * @param pf[IN] PageFile to write to
 * @return 0 if successful. Return an error code if there is an error.
 */
RC BTNonLeafNode::write(PageId pid, PageFile& pf) { 
  return pf.write(pid, buffer);
}

/*
 * Return the number of keys stored in the node.
 * @return the number of keys in the node
 */
int BTNonLeafNode::getKeyCount() {
  int* keyCount = (int*) &buffer[KEY_COUNT_INDEX];
  return *keyCount;
}


/*
 * Insert a (key, pid) pair to the node.
 * @param key[IN] the key to insert
 * @param pid[IN] the PageId to insert
 * @return 0 if successful. Return an error code if the node is full.
 */
RC BTNonLeafNode::insert(int key, PageId pid) {
  int keyCount = getKeyCount();
    // Check if the node is full
  if (keyCount == MAX_ELEMS)
    return RC_NODE_FULL;
    
    // Find position to insert new key
  int insertIndex = 0;
  int eid = 0;
  for (; eid < keyCount; eid++) {
    int* currentKey = (int*) &buffer[ELEM_SIZE * eid + sizeof(PageId)];
    if (*currentKey >= key)
      break;
  }
  insertIndex = ELEM_SIZE * eid + sizeof(PageId);
  
    // Shift existing pairs over
  memmove(&buffer[insertIndex + ELEM_SIZE], &buffer[insertIndex], KEY_COUNT_INDEX - (insertIndex + ELEM_SIZE));
  
    // Insert new pair
  int* newKey = (int*) &buffer[insertIndex];
  PageId* newPid = (PageId*) &buffer[insertIndex + sizeof(int)];
  *newKey = key;
  *newPid = pid;
  
    // Increase keyCount;
  int* keyCountBuf = (int*) &buffer[KEY_COUNT_INDEX];
  (*keyCountBuf)++;
  return 0; 
}

/*
 * Insert the (key, pid) pair to the node
 * and split the node half and half with sibling.
 * The middle key after the split is returned in midKey.
 * @param key[IN] the key to insert
 * @param pid[IN] the PageId to insert
 * @param sibling[IN] the sibling node to split with. This node MUST be empty when this function is called.
 * @param midKey[OUT] the key in the middle after the split. This key should be inserted to the parent node.
 * @return 0 if successful. Return an error code if there is an error.
 */
RC BTNonLeafNode::insertAndSplit(int key, PageId pid, BTNonLeafNode& sibling, int& midKey) { 
  int keyCount = getKeyCount();
    // Index of the key in the middle
  int splitCount = keyCount / 2;
  int splitOffset = 0;
  
    // Find position to insert new key
  int insertIndex = 0;
  int eid = 0;
  for (; eid < keyCount; eid++) {
    int* currentKey = (int*) &buffer[ELEM_SIZE * eid + sizeof(PageId)];
    if (*currentKey >= key)
      break;
  }
  insertIndex = ELEM_SIZE * eid + sizeof(PageId);
  
    // Check which node to insert new key
  if (insertIndex > splitCount)
    splitOffset = 1;
  
    // Insert second half into sibling node
  for (int i = splitCount + splitOffset; i < keyCount; i++) {
    int* keyToCopy = (int*) &buffer[i * ELEM_SIZE];
    PageId* pidToCopy = (PageId*) &buffer[i * ELEM_SIZE + sizeof(int)];
    sibling.insert(*keyToCopy, *pidToCopy);
  }  
    
    // Clear second half in current node
  memset(&buffer[(splitCount + splitOffset) * ELEM_SIZE], -1, (keyCount - splitCount - splitOffset) * ELEM_SIZE);
  
    // Insert new key into the correct node
  if (insertIndex > splitCount)
    sibling.insert(key, pid);
  else
    insert(key, pid);
  
    // Update current key count
  int* currentKeyCount = (int*) &buffer[KEY_COUNT_INDEX];
  *currentKeyCount = splitCount + 1;
  
    // Return middle key, which is always at the end of the first node
  midKey = *((int*) &buffer[sizeof(PageId) + *currentKeyCount * ELEM_SIZE]);
  PageId* midPid = (PageId*) &buffer[sizeof(PageId) + *currentKeyCount * ELEM_SIZE + sizeof(int)];
  sibling.setFirstPid(*midPid);
  memset(&buffer[sizeof(PageId) + *currentKeyCount * ELEM_SIZE], -1, ELEM_SIZE);
  *currentKeyCount--;
  
  return 0; 
}

/*
 * Given the searchKey, find the child-node pointer to follow and
 * output it in pid.
 * @param searchKey[IN] the searchKey that is being looked up.
 * @param pid[OUT] the pointer to the child node to follow.
 * @return 0 if successful. Return an error code if there is an error.
 */
RC BTNonLeafNode::locateChildPtr(int searchKey, PageId& pid) {
  int keyCount = getKeyCount();
  for (int i = 0; i < keyCount; i++) {
    int* currentKey = (int*) &buffer[i * ELEM_SIZE + sizeof(PageId)];
    if (searchKey < *currentKey) {
      pid = *((PageId*) &buffer[i * ELEM_SIZE]);
      return 0;
    }
  }
  pid = *((PageId*) &buffer[keyCount * ELEM_SIZE]);
  return 0; 
}

/*
 * Initialize the root node with (pid1, key, pid2).
 * @param pid1[IN] the first PageId to insert
 * @param key[IN] the key that should be inserted between the two PageIds
 * @param pid2[IN] the PageId to insert behind the key
 * @return 0 if successful. Return an error code if there is an error.
 */
RC BTNonLeafNode::initializeRoot(PageId pid1, int key, PageId pid2) {
  PageId* firstPid = (PageId*) &buffer[0];
  *firstPid = pid1;
  insert(key, pid2);
  return 0; 
}
